import requests,json
from support.utils.toolkit import Log

class Localai():
    LOCALAI_SERVER_URL="http://localai:8080"
    EMBEDDING_URI="v1/embeddings"
    DEFAULT_EMBED_MODEL = "bert-embeddings"
    DEFAULT_MAX_TOKEN = 512
    DEFAULT_HEADERS = {
        'accept': 'application/json',
        'Content-Type': 'application/json',
        'User-Agent': 'MY RGA Agent 1.0'
    }

    def embeddings(self, text, *, model=None):
        ret = None
        if model is None:
            model=Localai.DEFAULT_EMBED_MODEL
        data = {
            "model": model,
            "max_tokens": Localai.DEFAULT_MAX_TOKEN,
            "input": text
        }
        result = requests.post("{}/{}".format(Localai.LOCALAI_SERVER_URL, Localai.EMBEDDING_URI), headers=Localai.DEFAULT_HEADERS, data=json.dumps(data))
        if(result.status_code == 200):
            jdata = json.loads(result.text)
            ret =jdata['data'][0]['embedding']
        else:
            Log.e(result.text)
        return ret

